#include "chrono"
#include "opencv2/opencv.hpp"
#include <algorithm>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#include <string>
#include <vector>
#include <iostream>
#include <cmath>
#include <float.h>
#include <iomanip>

using namespace std;

void test_ball_in_kuang(const cv::Mat& image)
{
    if (image.empty()) {
        std::cout << "图像为空！" << std::endl;
        return;
    }
    
    // 提取中间1/3区域
    int height = image.rows;
    int width = image.cols;
    int roi_width= width / 3;
    int roi_x = width / 3; // 从图像1/3处开始

    cv::Rect roi(roi_x, 0, roi_width, height);
    cv::Mat roi_image = image(roi).clone();
    
    std::cout << "ROI区域大小: " << roi_image.cols << "x" << roi_image.rows << std::endl;
    
    // 定义5种颜色的BGR参考值
    std::vector<cv::Vec3b> color_refs = {
        cv::Vec3b(20, 15, 100),   // 红色小球 (BGR)
        cv::Vec3b(60, 150, 150),  // 黄色小球 (BGR)
        cv::Vec3b(10, 10, 10),    // 黑色小球 (BGR)
        cv::Vec3b(90, 45, 13),    // 蓝色小球 (BGR)
        cv::Vec3b(120, 120, 120)  // 灰白色 (BGR)
    };

    // 定义5种颜色的BGR查找表
    std::vector<cv::Vec3b> color_lut = {
        cv::Vec3b(0, 0, 255),   // 红色小球 (BGR)
        cv::Vec3b(0, 255, 255),  // 黄色小球 (BGR)
        cv::Vec3b(0, 0, 0),    // 黑色小球 (BGR)
        cv::Vec3b(255, 0, 0),    // 蓝色小球 (BGR)
        cv::Vec3b(255, 255, 255)  // 灰白色 (BGR)
    };
    
    std::vector<std::string> color_names = {"红色", "黄色", "黑色", "蓝色", "灰白色"};
    
    // 创建颜色分类结果图像
    cv::Mat classified_image = cv::Mat::zeros(roi_image.size(), CV_8UC3);
    cv::Mat color_mask = cv::Mat::zeros(roi_image.size(), CV_8UC1);
    
    // 使用欧式距离进行颜色分类
    for (int y = 0; y < roi_image.rows; y++) {
        for (int x = 0; x < roi_image.cols; x++) {
            cv::Vec3b pixel = roi_image.at<cv::Vec3b>(y, x);
            
            double min_distance = DBL_MAX;
            int best_class = 0;
            
            // 计算与每种参考颜色的欧式距离
            for (int i = 0; i < color_refs.size(); i++) {
                cv::Vec3b ref_color = color_refs[i];
                double distance = sqrt(pow(pixel[0] - ref_color[0], 2) + 
                                     pow(pixel[1] - ref_color[1], 2) + 
                                     pow(pixel[2] - ref_color[2], 2));
                
                if (distance < min_distance) {
                    min_distance = distance;
                    best_class = i;
                }
            }
            
            // 分配分类颜色
            classified_image.at<cv::Vec3b>(y, x) = color_lut[best_class];
            color_mask.at<uchar>(y, x) = best_class;
        }
    }
    
    // 统计每种颜色的像素数量
    std::vector<int> color_counts(3, 0);
    for (int y = 0; y < color_mask.rows; y++) {
        for (int x = 0; x < color_mask.cols; x++) {
            color_counts[color_mask.at<uchar>(y, x)]++;
        }
    }
    
    std::cout << "\n颜色统计结果:" << std::endl;
    for (int i = 0; i < color_names.size(); i++) {
        std::cout << color_names[i] << ": " << color_counts[i] << " 像素" << std::endl;
    }
    
    // 基于颜色分类结果的小球检测
    cv::Mat result_image = roi_image.clone();
    cv::Mat all_contours_image = cv::Mat::zeros(roi_image.size(), CV_8UC3);
    
    // 将颜色分类结果转换为灰度图
    cv::Mat classified_gray;
    cv::cvtColor(classified_image, classified_gray, cv::COLOR_BGR2GRAY);
    
    // 计算图像面积的10%作为面积阈值
    int image_area = roi_image.rows * roi_image.cols;
    int min_area_threshold = image_area * 0.1;
    
    std::cout << "\n图像总面积: " << image_area << " 像素" << std::endl;
    std::cout << "面积阈值(10%): " << min_area_threshold << " 像素" << std::endl;
    
    std::vector<int> ball_counts_by_color(4, 0); // 只统计4种小球颜色
    std::vector<std::string> ball_names = {"红色小球", "黄色小球", "黑色小球", "蓝色小球"};
    std::vector<cv::Scalar> display_colors = {
        cv::Scalar(0, 0, 255),    // 红色显示
        cv::Scalar(0, 255, 255),  // 黄色显示
        cv::Scalar(128, 128, 128), // 黑色显示
        cv::Scalar(255, 0, 0)     // 蓝色显示
    };
    
    std::cout << "\n基于灰度化轮廓检测的结果:" << std::endl;
    
    // 对每种颜色分别进行检测（排除灰白色背景）
    for (int color_idx = 0; color_idx < 4; color_idx++) {
        // 创建当前颜色的二值掩码
        cv::Mat color_binary = cv::Mat::zeros(color_mask.size(), CV_8UC1);
        for (int y = 0; y < color_mask.rows; y++) {
            for (int x = 0; x < color_mask.cols; x++) {
                if (color_mask.at<uchar>(y, x) == color_idx) {
                    color_binary.at<uchar>(y, x) = 255;
                }
            }
        }
        
        // 跳过像素数量太少的颜色
        if (color_counts[color_idx] < 100) continue;
        
        // 检测轮廓
        std::vector<std::vector<cv::Point>> contours;
        std::vector<cv::Vec4i> hierarchy;
        cv::findContours(color_binary, contours, hierarchy, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_SIMPLE);
        
        std::cout << "\n" << color_names[color_idx] << "颜色检测:" << std::endl;
        
        for (size_t i = 0; i < contours.size(); i++) {
            double area = cv::contourArea(contours[i]);
            
            // 排除面积小于图像面积10%的轮廓
            if (area < min_area_threshold) {
                std::cout << "  轮廓 " << (i+1) << ": 面积=" << (int)area << " (小于阈值，跳过)" << std::endl;
                continue;
            }
            
            // 计算轮廓的圆形度
            double perimeter = cv::arcLength(contours[i], true);
            double circularity = 4 * CV_PI * area / (perimeter * perimeter);
            
            // 计算边界矩形和长宽比
            cv::Rect boundRect = cv::boundingRect(contours[i]);
            double aspect_ratio = (double)boundRect.width / boundRect.height;
            
            // 计算最小外接圆
            cv::Point2f center;
            float radius;
            cv::minEnclosingCircle(contours[i], center, radius);
            
            // 进行圆度和长宽比筛选
            // if (circularity > 0.7 && aspect_ratio > 0.9 && aspect_ratio < 1.3) {
            if (circularity > 0.2 && aspect_ratio > 0.5 && aspect_ratio < 1.8) {
                // 绘制检测结果
                cv::drawContours(result_image, contours, i, display_colors[color_idx], 2);
                cv::drawContours(all_contours_image, contours, i, display_colors[color_idx], -1);
                
                // 标注小球中心和信息
                cv::circle(result_image, center, 3, cv::Scalar(0, 255, 0), -1);
                cv::circle(result_image, center, radius, display_colors[color_idx], 2);
                
                // 添加文字标注
                cv::putText(result_image, ball_names[color_idx], 
                           cv::Point(boundRect.x, boundRect.y - 10), 
                           cv::FONT_HERSHEY_SIMPLEX, 0.5, display_colors[color_idx], 2);
                
                ball_counts_by_color[color_idx]++;
                
                std::cout << "  ✓ " << ball_names[color_idx] << " " << ball_counts_by_color[color_idx] 
                         << ": 面积=" << (int)area 
                         << ", 圆形度=" << std::fixed << std::setprecision(3) << circularity
                         << ", 长宽比=" << std::fixed << std::setprecision(3) << aspect_ratio
                         << ", 半径=" << std::fixed << std::setprecision(1) << radius
                         << ", 中心(" << (int)center.x << "," << (int)center.y << ")" << std::endl;
            } else {
                std::cout << "  ✗ 不符合条件: 面积=" << (int)area 
                         << ", 圆形度=" << std::fixed << std::setprecision(3) << circularity
                         << ", 长宽比=" << std::fixed << std::setprecision(3) << aspect_ratio << std::endl;
            }
        }
    }
    
    // 统计总的小球检测结果
    int total_balls = 0;
    std::cout << "\n=== 色块检测汇总 ===" << std::endl;
    for (int i = 0; i < 4; i++) {
        if (ball_counts_by_color[i] > 0) {
            std::cout << ball_names[i] << ": " << ball_counts_by_color[i] << " 个" << std::endl;
            total_balls += ball_counts_by_color[i];
        }
    }
    std::cout << "总计检测到色块: " << total_balls << " 个" << std::endl;
    std::cout << "检测条件: 面积>图像10%, 圆形度>0.7, 长宽比0.8-1.25" << std::endl;
    
    // 显示结果
    cv::imshow("原始ROI区域", roi_image);
    cv::imshow("颜色分类结果", classified_image);
    cv::imshow("颜色分类灰度图", classified_gray);
    cv::imshow("轮廓检测结果", result_image);
    cv::imshow("所有检测到的色块轮廓", all_contours_image);
    
    // 保存结果
    cv::imwrite("/tmp/roi_original.jpg", roi_image);
    cv::imwrite("/tmp/color_classified.jpg", classified_image);
    cv::imwrite("/tmp/classified_gray.jpg", classified_gray);
    cv::imwrite("/tmp/contour_detection_result.jpg", result_image);
    cv::imwrite("/tmp/all_color_contours.jpg", all_contours_image);
    
    std::cout << "\n结果已保存到 /tmp/ 目录" << std::endl;
    std::cout << "按任意键继续..." << std::endl;
    cv::waitKey(0);
    cv::destroyAllWindows();
}

int main(int argc, char** argv)
{
    string photo_file_path = "/home/firefly/1001/b.jpg";
    if(argc == 2) {
        photo_file_path = argv[1];
    }
    cv::Mat imagek = cv::imread(photo_file_path);
    cv::resize(imagek, imagek, cv::Size(640, 360));
    test_ball_in_kuang(imagek);
    
    return 0;
}
